# 处理图片

## 1. 分patch 把一张图片 (1, 28, 28) -> (1, 49, 16) 49个片，每片16个像素点

## 2. 合成完整图像 transformer 模型输出 (1, 49, 4, 4) 合成完整图像

import numpy as np 

class ImageHanle:
    def  __init__(self, patch_w, patch_h):
        # 保证可以整除
        self.patch_w = patch_w
        self.patch_h = patch_h 

    def patch_image(self, single_img):
        # single_img: (channel, W, H)   
        
        patchs = []
        if len(single_img.shape) == 3:
            channel, W, H = single_img.shape
            self.W = W 
            self.H = H 
            self.patch_num_w = W // self.patch_w
            self.patch_num_h = H // self.patch_h 

            patchs = []
            for ww in range(self.patch_num_w):
                for hh in range(self.patch_num_h):
                    patchs.append(single_img[:, ww * self.patch_w:ww * self.patch_w + self.patch_w, hh * self.patch_h:hh * self.patch_h + self.patch_h].reshape(-1))
        elif len(single_img.shape) == 2:
            W, H = single_img.shape
            self.W = W 
            self.H = H 
            self.patch_num_w = W // self.patch_w
            self.patch_num_h = H // self.patch_h 
            for ww in range(self.patch_num_w):
                for hh in range(self.patch_num_h):
                    patchs.append(single_img[ww * self.patch_w:ww * self.patch_w + self.patch_w, hh * self.patch_h:hh * self.patch_h + self.patch_h].reshape(-1))

        return np.stack(patchs, axis=0)
    
    def cons_img(self, img_patchs, channels):
        # 组合图片 img_patches: (len, size)
        img = np.zeros((channels, self.W, self.H))
        patch_num, size = img_patchs.shape
        index = 0
        for ww in range(self.patch_num_w):
                for hh in range(self.patch_num_h):
                   img[:, ww * self.patch_w:ww * self.patch_w + self.patch_w, hh * self.patch_h:hh * self.patch_h + self.patch_h] = img_patchs[index].reshape((channels, self.patch_w, self.patch_h))
                   index += 1
        return img 



if __name__ == "__main__":
    a1 = np.random.rand(3, 50, 60)
    print(a1.shape)
    
    image_hanle = ImageHanle(10, 12)
    # print(image_hanle.patch_image(a1).shape)
    img_patches = image_hanle.patch_image(a1)
    print(img_patches.shape)

    img = image_hanle.cons_img(img_patches, 3)
    print(img.shape)
    # img = img.squeeze(0)

    print(img - a1)